# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/11 16:06 
@Author  : ZhangShenao 
@File    : 2.输出解析.py 
@Desc    : OutParser输出解析
"""
import dotenv
import pandas as pd
from langchain.output_parsers import ResponseSchema, StructuredOutputParser
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.prompts import ChatPromptTemplate

dotenv.load_dotenv()

# 创建原始模板
template = '''你是一位专业的鲜花店文案撰写员\n
对于售价为 {price} 元的 {flower_name}，请帮我生成一个简短且吸引人的宣传文案。
{format_instructions}
'''

# 创建模型实例
llm = ChatZhipuAI(model_name="glm-4-air")

# 定义想要接收的响应格式
response_schemas = [
    ResponseSchema(name='description', description='鲜花的描述文案', type='string'),
    ResponseSchema(name='reason', description='为什么要这样写这个文案', type='string')
]

# 创建输出解析器
# 使用StructuredOutputParser,可以将LLM的输出解析为指定的结构化的格式
parser = StructuredOutputParser.from_response_schemas(response_schemas)

# 生成格式指令
format_instructions = parser.get_format_instructions()

# 预填充格式指令,生成Prompt
# Prompt中可以嵌入对输出格式的定义
prompt = ChatPromptTemplate.from_template(template).partial(format_instructions=format_instructions)

# 构造Chain
chain = prompt | llm | parser

# 创建一个空的DataFrame用于存储结果
df = pd.DataFrame(columns=["flower", "price", "description", "reason"])  # 先声明列名

# 构造输入数据
prices = ['50', '30', '20']
flowers = ['玫瑰花', '百合花', '康乃馨']

# 调用Chain,处理输出结果
for price, flower in zip(prices, flowers):
    output = chain.invoke({"price": price, "flower_name": flower})

    # 在解析后的结果中添加flower和price字段
    output['flower'] = flower
    output['price'] = price

    # 将解析后的输出添加到DataFrame中
    df.loc[len(df)] = output

# 打印结果字典
print(df.to_dict(orient='records'))

# 保存DataFrame到CSV文件
df.to_csv("./flowers_with_descriptions.csv", index=False)
